#!/usr/bin/env python3

#-------------------------------------------------------------------------------
# Mongoose/Tests/runTests
#-------------------------------------------------------------------------------

# Mongoose Graph Partitioning Library, Copyright (C) 2017-2023,
# Scott P. Kolodziej, Nuri S. Yeralan, Timothy A. Davis, William W. Hager
# Mongoose is licensed under Version 3 of the GNU General Public License.
# Mongoose is also available under other licenses; contact authors for details.
# SPDX-License-Identifier: GPL-3.0-only

#-------------------------------------------------------------------------------

from subprocess import call, check_output
import os       # For filesystem access
import sys      # For sys.exit()
import argparse # For parsing command-line arguments
import urllib.request, urllib.parse, urllib.error   # For downloading the ssget index
import ssl
import tarfile  # For un-tar/unzipping matrix files
import csv      # For reading the ssget index
import shutil   # For using 'which'
import platform

#-------------------------------------------------------------------------------
# main: parse arguments, run the tests, and determine test coverage
#-------------------------------------------------------------------------------

def main():
    print(platform.sys.version)

    # Parse the command-line arguments
    args = parseArguments()

    # Run tests if needed
    if args.tests != "none":
        runTests(args)

    # If the coverage flag is on, run gcov (or similar)
    if args.coverage or args.html_coverage:
        runCoverageUtility(args)

#-------------------------------------------------------------------------------
# runTests: run the Mongoose tests
#-------------------------------------------------------------------------------

def runTests(args):
    # Create or locate matrix temporary storage directory
    matrix_dir = getMatrixDirectory()
    print("Matrix directory: " + matrix_dir)

    # find the directory with the compiled tests
    Mongoose_tests_dir = checkLocation_tests_dir()

    # Download the matrix stats csv file
    stats_file = downloadStatsFile(matrix_dir)

    with open(stats_file, 'r') as f:
        reader = csv.reader(f)

        # Matrix IDs are not listed in the stats file - we just have to keep count
        matrix_id = 0
        for row in reader:

            if len(row) == 13: # Only rows with 13 elements represent matrix data
                matrix_id += 1

                # Check if the matrix ID is in the proper range and
                # that the matrix is real and symmetric
                isInBounds = ((matrix_id >= args.id_min) and (matrix_id <= args.id_max))
                isSquare = (row[2] == row[3])
                isReal = (row[5] == '1')

                if (isInBounds and isSquare and isReal):
                    if args.ids is None or matrix_id in args.ids:
                        matrix_name = row[0] + '/' + row[1] + '.tar.gz'
                        gzip_path = matrix_dir + row[0] + '_' + row[1] + '.tar.gz'
                        matrix_path = matrix_dir + row[1] + '/' + row[1] + ".mtx"
                        print("matrix_name: " + matrix_name)
                        print("gzip_path:   " + gzip_path)
                        print("matrix_path: " + matrix_path)

                        matrix_exists = os.path.isfile(gzip_path)
                        if matrix_exists:
                            print("matrix exists at gzip_path")
                            tar = tarfile.open(gzip_path, mode='r:gz')
                            matrix_files = tar.getnames()
                            print(matrix_files)
                        else:
                            # Download matrix if it doesn't exist
                            try:
                                print("Downloading " + matrix_name)
                                url = "https://sparse.tamu.edu/MM/" + matrix_name
                                print("url: " + url)
                                with urllib.request.urlopen(url) as response, open(gzip_path, 'wb') as out_file:
                                    shutil.copyfileobj(response, out_file)
                            except:
                                print("Downloading " + matrix_name + "via HTTPS failed. Falling back to HTTP...")
                                url = "http://sparse.tamu.edu/MM/" + matrix_name
                                print("url: " + url)
                                with urllib.request.urlopen(url) as response, open(gzip_path, 'wb') as out_file:
                                    shutil.copyfileobj(response, out_file)
                            tar = tarfile.open(gzip_path, mode='r:gz')
                            tar.extractall(path=matrix_dir) # Extract the matrix from the tar.gz file
                            tar.close()

                        # Determine which test executables to run
                        if args.tests == 'all':
                            print("Calling ALL Tests...")
                        if args.tests == 'memory' or args.tests == 'all':
                            print("Calling Memory Test...")
                            status = call([Mongoose_tests_dir + "mongoose_test_memory", matrix_path])
                            if status:
                                print("Error! Memory Test Failure")
                                # cleanup(args, row, matrix_exists, gzip_path)
                                sys.exit(status)
                        if args.tests == 'valgrind' or args.tests == 'all':
                            print("Calling Valgrind Test...")
                            # valgrind = find_executable('valgrind')
                            valgrind = shutil.which('valgrind')
                            if valgrind:
                                status = call([valgrind + " --leak-check=full " + Mongoose_tests_dir + "mongoose_test_memory", matrix_path])
                                if status:
                                    print("Error! Valgrind Test Failure")
                                    # cleanup(args, row, matrix_exists, gzip_path)
                                    sys.exit(status)
                            else:
                                print("\033[91mERROR!\033[0m Unable to find Valgrind. Skipping Valgrind Test...")
                        if args.tests == 'io' or args.tests == 'all':
                            print("Calling I/O Test...")
                            status = call([Mongoose_tests_dir + "mongoose_test_io", matrix_path, "1"])
                            if status:
                                print("Error! I/O Test Failure")
                                # cleanup(args, row, matrix_exists, gzip_path)
                                sys.exit(status)
                        if args.tests == 'edgesep' or args.tests == 'all':
                            print("Calling Edge Separator Test...")
                            target_split = args.target_split
                            status = call([Mongoose_tests_dir + "mongoose_test_edgesep", matrix_path, str(target_split)])
                            if status:
                                print("Error! Edge Separator Test Failure")
                                # cleanup(args, row, matrix_exists, gzip_path)
                                sys.exit(status)
                        if args.tests == 'performance' or args.tests == 'all':
                            print("Calling Performance Test...")
                            status = call([Mongoose_tests_dir + "mongoose_test_performance", matrix_path, row[0] + '_' + row[1] + '_performance.txt'])
                            if status:
                                print("Error! Performance Test Failure")
                                # cleanup(args, row, matrix_exists, gzip_path)
                                sys.exit(status)

                        # Delete the matrix only if we downloaded it and the keep
                        # flag is off
                        # cleanup(args, row, matrix_exists, gzip_path)

    # delete the ssstats.csv file
    os.remove(stats_file)

#-------------------------------------------------------------------------------
# cleanup: remove matrices if downloaded and keep flag is off
#-------------------------------------------------------------------------------

def cleanup(args, row, matrix_exists, gzip_path):
    print("cleanup")
    if args.purge or not (args.keep or matrix_exists):
        matrix_dir = getMatrixDirectory()
        print("cleanup matrix_dir: " + matrix_dir)
        files = os.listdir(matrix_dir + row[1])
        print(files)
        for file in files:
            print("remove file: " + file)
            os.remove(os.path.join(matrix_dir + row[1] + '/', file))
        print("remove dir: " + matrix_dir + row[1])
        os.rmdir(matrix_dir + row[1])
        print("remove gzip_path: " + gzip_path)
        os.remove(gzip_path)

#-------------------------------------------------------------------------------
# getMatrixDirectory: check if the Matrix directory exists; if not create it
#-------------------------------------------------------------------------------

def getMatrixDirectory():
    matrix_dir = "./Matrix/"
    if (not os.path.exists(matrix_dir)):
        os.makedirs(matrix_dir)
    return matrix_dir

#-------------------------------------------------------------------------------
# downloadStatsFile: download the sparse.tamu.edu/files/ssstats.csv file
#-------------------------------------------------------------------------------

def downloadStatsFile(matrix_dir):
    stats_file = "ssstats.csv"
    print("downloading: ssstats.csv")
    url = "http://sparse.tamu.edu/files/ssstats.csv"
    with urllib.request.urlopen(url) as response, open(stats_file, 'wb') as out_file:
        shutil.copyfileobj(response, out_file)
    return stats_file

#-------------------------------------------------------------------------------
# runCoverageUtility: run gcov
#-------------------------------------------------------------------------------

def runCoverageUtility(args):
    if args.gcov:
        gcov = args.gcov
    else:
        # gcov = find_executable('gcov')
        gcov = shutil.which('gcov')

    # find the Mongoose/Source directory
    Mongoose_source_dir = checkLocation_source_dir()

    # find the directory with compiled *.o files
    Mongoose_object_dir = checkLocation_object_dir()

    if (len(Mongoose_object_dir) > 0 and len (Mongoose_source_dir) > 0):

        if gcov:
            # Determine if we are using GCC gcov or LLVM gcov
            gcov_version = check_output([gcov, "--version"])
            if gcov_version.find('LLVM') == -1:
                call([gcov + " -o " + Mongoose_object_dir + " " + Mongoose_source_dir + "/*.cpp"], shell=True)
            else:
                call(gcov + " -o=" + Mongoose_object_dir + " " + Mongoose_source_dir + "/*.cpp", shell=True)

        # gcovr = find_executable('gcovr')
        gcovr = shutil.which('gcovr')
        if gcovr:
            if args.html_coverage:
                print("Running gcovr with HTML generation")
                call([gcovr,
                      "--html",
                      "--html-details",
                      "--output=coverage.html",
                      "--gcov-executable=" + gcov,
                      "--object-directory=" + Mongoose_object_dir,
                      "--root=" + Mongoose_source_dir])
            else:
                print("Running gcovr without HTML generation")
                call([gcovr, "--gcov-executable=" + gcov, "--object-directory=" + Mongoose_object_dir, "--root=" + Mongoose_source_dir])
        else:
            print("\033[91mERROR!\033[0m Cannot generate HTML coverage report, gcovr not found!")

#-------------------------------------------------------------------------------
# parseArguments: parse input arguments
#-------------------------------------------------------------------------------

def parseArguments():
    parser = argparse.ArgumentParser(
                        description='Run tests on the Mongoose library.')
    parser.add_argument('-k', '--keep',
                        action='store_true',
                        help='do not remove downloaded files when test is complete')
    parser.add_argument('-p', '--purge',
                        action='store_true',
                        help='force remove downloaded matrix files when complete')
    parser.add_argument('-min',
                        action='store',
                        metavar='min_id',
                        type=int,
                        default=1,
                        dest='id_min',
                        help='minimum matrix ID to run tests on [default: 1]')
    parser.add_argument('-max',
                        action='store',
                        metavar='max_id',
                        type=int,
                        default=2757,
                        dest='id_max',
                        help='maximum matrix ID to run tests on [default: 2757]')
    parser.add_argument('-i', '--ids',
                        action='store',
                        nargs='+',
                        metavar='matrix_ID',
                        type=int,
                        help='list of matrix IDs to run tests on')
    parser.add_argument('-t', '--tests',
                        choices=['all', 'memory', 'io', 'edgesep', 'performance', 'valgrind', 'none'],
                        default='none',
                        help='choice of which tests to run')
    parser.add_argument('-s', '--split',
                        action='store',
                        metavar='target_split',
                        type=float,
                        dest='target_split',
                        default=0.5,
                        help='target split ratio for edge separator test [default: 0.5]')
    parser.add_argument('-c', '--coverage',
                        action='store_true',
                        help='generate coverage information')
    parser.add_argument('--html-coverage',
                        action='store_true',
                        help='generate html coverage pages if gcovr is available')
    parser.add_argument('--gcov',
                        action='store',
                        metavar='gcov_path',
                        help='path to gcov tool')

    return parser.parse_args()

#-------------------------------------------------------------------------------
# checkLocation_object_dir: look for the compiled Mongoose object files
#-------------------------------------------------------------------------------

def checkLocation_object_dir():
    object_dir1 = "./CMakeFiles/Mongoose_static_dbg.dir/Source"
    object_dir2 = "./Mongoose/CMakeFiles/Mongoose_static_dbg.dir/Source"
    if (os.path.isdir(object_dir1)):
        # built in SuiteSparse/Mongoose/build (stand-alone)
        object_dir = object_dir1
    elif (os.path.isdir(object_dir2)):
        # built in SuiteSparse/build (via the root SuiteSparse/CMakeLists.txt)
        object_dir = object_dir2
    else:
        # no object directory
        object_dir = ""
    print("object directory: " + object_dir)
    return object_dir

#-------------------------------------------------------------------------------
# checkLocation_source_dir: look for the Mongoose/Source files
#-------------------------------------------------------------------------------

def checkLocation_source_dir():
    source_dir1 = "../Source"
    source_dir2 = "../Mongoose/Source"
    if (os.path.isdir(source_dir1)):
        # built in SuiteSparse/Mongoose/build (stand-alone)
        source_dir = source_dir1
    elif (os.path.isdir(source_dir2)):
        # built in SuiteSparse/build (via the root SuiteSparse/CMakeLists.txt)
        source_dir = source_dir2
    else:
        source_dir = ""
    print("Mongoose/Source directory: " + source_dir)
    return source_dir

#-------------------------------------------------------------------------------
# checkLocation_tests_dir: look for the compiled Mongoose/tests
#-------------------------------------------------------------------------------

def checkLocation_tests_dir():
    tests_dir = ""
    tests_dir1 = "./tests/"
    tests_dir2 = "../Mongoose/tests/"
    # List of multi-config configurations (e.g., for MSVC)
    # FIXME: Add more multi-config configurations to list?
    # FIXME: Instead of blindly looking for existing directories, check which
    #        configuration was selected in ctest flag. E.g., `ctest -C Release`
    config_dirs = ["Release/", "Debug/", ""]
    for config_dir in config_dirs:
        print("checking for: " + tests_dir1 + config_dir)
        if (os.path.isdir(tests_dir1 + config_dir)):
            # built in SuiteSparse/Mongoose/build (stand-alone)
            tests_dir = tests_dir1 + config_dir
            break
        elif (os.path.isdir(tests_dir2 + config_dir)):
            # built in SuiteSparse/build (via the root SuiteSparse/CMakeLists.txt)
            tests_dir = tests_dir2 + config_dir
            break
    if not tests_dir:
        print(
            "\n\033[91mERROR!\033[0m Looks like you might not be running this from "
            "your build directory.\n\n"
            "Make sure that... \n\n"
            "  * You are in your build directory (e.g. Mongoose/build) and\n"
            "  * You have built Mongoose ('cmake ..' followed by 'make')\n")
        sys.exit('Test failed: cannot find Mongoose')
    print("compiled tests directory: " + tests_dir)
    return tests_dir

#-------------------------------------------------------------------------------
# call the main method
#-------------------------------------------------------------------------------

if __name__=="__main__":
    main()

